use crate::block::Block;
use crate::consts::{BIN_SIZE, SEGMENT_MASK, SEGMENT_SIZE, SMALL_PAGE_SHIFT, SMALL_SIZE_MAX};
use crate::page::{Page, PageKind};
use crate::segment::{Segment, SegmentGroup, SegmentOwner};
use crate::utils::access_guard::AccessGuard;
use crate::utils::atomic_single_linked_list::AtomicSingleLinkedList;
use crate::utils::bin_from_size;
use crate::utils::double_linked_list::DoubleLinkedList;
use crate::utils::os_box::{OsBox, OsBoxBody};
use std::alloc::Layout;
use std::array;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;
#[cfg(not(feature = "multi-thread"))]
use std::sync::atomic::Ordering;

thread_local! {
    static THREAD_LOCAL_HEAP: Heap = Heap::new();
}

#[cfg(not(feature = "multi-thread"))]
static SINGLE_THREAD_SINGLETON_HEAP: std::sync::atomic::AtomicPtr<HeapInner> =
    std::sync::atomic::AtomicPtr::new(std::ptr::null_mut());

pub(crate) struct HeapInner {
    // active_pages 中页面一定是有可分配的 Block 的
    // 触发 active_pages -> sleep_pages 操作：
    //   pages 的 malloc 接口返回 None，pages 上没有任何空间可供分配。
    active_pages: [DoubleLinkedList<Page>; BIN_SIZE],
    // sleep_pages
    // 触发 sleep_pages -> active_pages 操作：
    //   local free 一个 sleep page 中的 Block，直接将对应 page 唤醒到 active_pages
    //
    // 如果一个 page 成为 sleep page 后仅进行过 thread free，一次 local free 调用都没有过，
    // 他将成为存在可用空间的 sleep page，这种情况下，page 会添加到 delayed_week_up_pages，
    // 然后由 heap 选择时机统一回收。
    sleep_pages: DoubleLinkedList<Page>,
    // 由其他线程释放出空间的 Page，暂时未添加到 active_pages
    delayed_week_up_pages: AtomicSingleLinkedList<Page>,

    // 是当前线程所有 segment 的 Owner
    segment_group: SegmentGroup,
}

impl OsBoxBody for HeapInner {
    fn init() -> Self {
        Self {
            active_pages: array::from_fn(|_| DoubleLinkedList::default()),
            sleep_pages: Default::default(),
            delayed_week_up_pages: Default::default(),
            segment_group: Default::default(),
        }
    }

    fn before_drop(&mut self) {
        // TODO 进行一系列检查
        // 逃逸检查，heap 释放但是seg仍然在使用的情况
        // 对于仍然在使用的页面和堆，标记并另存
    }
}

impl Heap {
    fn new() -> Self {
        Heap(OsBox::new().expect("Failed to alloc thread local data: Out of memory"))
    }

    fn week_up_delayed_pages(&self) {
        let pages = self.delayed_week_up_pages.take();
        pages.iter(self).for_each(|page_ptr| {
            let page = self.access(page_ptr);
            page.thread_collect();
            if page.is_empty() {
                let segment_ptr = Self::ptr_segment(page_ptr.cast().as_ptr())
                    .expect("Never happened, the page ptr will not be null");
                let segment = self.access(segment_ptr);
                segment.page_used -= 1;
            }
            if page.is_in_sleep() {
                page.set_in_sleep(false);
                self.sleep_pages.leave(self, page_ptr);
                self.active_pages[page.bin()].push(self, page_ptr);
            }
            page.set_in_delayed(false);
        });
    }

    fn expend_segment(&self, bin: usize) {
        let (seg, free_pages) = match Segment::new(0, PageKind::Small, bin, SMALL_PAGE_SHIFT) {
            None => return,
            Some(v) => v,
        };
        self.active_pages[bin].join_tail(self, free_pages);
        self.segment_group.push(seg);
    }

    fn malloc_small(&self, layout: Layout) -> Option<NonNull<Block>> {
        debug_assert!(layout.size() <= SMALL_SIZE_MAX);
        let bin = bin_from_size(layout.size());
        let page_queue = &self.active_pages[bin];
        let page_ptr = page_queue
            .first()
            .or_else(|| {
                self.week_up_delayed_pages();
                page_queue.first()
            })
            .or_else(|| {
                self.expend_segment(bin);
                page_queue.first()
            });
        let page_ptr = match page_ptr {
            None => return None,
            Some(ptr) => ptr,
        };
        let page = self.access(page_ptr);
        if page.is_empty() {
            let segment_ptr = Self::ptr_segment(page_ptr.cast().as_ptr())
                .expect("Never happened, the page ptr will not be null");
            let segment = self.access(segment_ptr);
            segment.page_used += 1;
        }
        let ptr = page.alloc(layout);
        if page.is_full() {
            self.active_pages[bin].leave(self, page_ptr);
            self.sleep_pages.push(self, page_ptr);
            page.set_in_sleep(true);
        }
        ptr
    }

    #[inline(always)]
    fn ptr_segment(ptr: *mut u8) -> Option<NonNull<Segment>> {
        NonNull::new((ptr as usize & !SEGMENT_MASK) as *mut Segment)
    }

    fn drop_segment(&self, mut segment: SegmentOwner) {
        let segment = &mut segment.0;
        debug_assert_eq!(segment.page_used, 0);
        (0..segment.capacity())
            .map(|i| segment.page_at(i))
            .for_each(|page_ptr| {
                let page = self.access(page_ptr);
                debug_assert!(page.is_empty());
                debug_assert!(!page.is_in_delayed());
                debug_assert!(!page.is_in_sleep());
                self.active_pages[page.bin()].leave(self, page_ptr);
            });
    }

    #[cfg(feature = "multi-thread")]
    pub(crate) fn thread_local_with<F, R>(f: F) -> R
    where
        F: FnOnce(&Heap) -> R,
    {
        THREAD_LOCAL_HEAP.with(f)
    }

    #[cfg(not(feature = "multi-thread"))]
    pub(crate) fn single_thread_singleton() -> Heap {
        unsafe {
            if SINGLE_THREAD_SINGLETON_HEAP
                .load(Ordering::Relaxed)
                .is_null()
            {
                let heap = Heap::new();
                let ptr = heap.0.as_ptr();
                heap.0.leak();
                SINGLE_THREAD_SINGLETON_HEAP.store(ptr.as_ptr(), Ordering::SeqCst);
            }
            let ptr = SINGLE_THREAD_SINGLETON_HEAP.load(Ordering::Relaxed);
            Heap(OsBox::leak_box(NonNull::new_unchecked(ptr)))
        }
    }

    pub(crate) fn alloc(&self, layout: Layout) -> Option<NonNull<Block>> {
        let size = layout.pad_to_align().size();
        if size == 0 {
            return None;
        }
        if size <= SMALL_SIZE_MAX {
            self.malloc_small(layout)
        } else {
            panic!("not support non-small malloc for now")
        }
    }

    pub(crate) fn dealloc(&self, block_ptr: NonNull<Block>) {
        let segment_ptr = match Self::ptr_segment(block_ptr.as_ptr() as *mut u8) {
            None => return,
            Some(ptr) => ptr,
        };
        let segment = self.access(segment_ptr);
        let local = segment.is_local();
        let page_ptr = segment.ptr_page(block_ptr.as_ptr() as *mut u8);
        let page = segment.access(page_ptr);
        page.dealloc(block_ptr, local);
        if local {
            if page.is_in_sleep() {
                page.set_in_sleep(false);
                self.sleep_pages.leave(self, page_ptr);
                self.active_pages[page.bin()].push(self, page_ptr);
            }
            if page.is_empty() {
                segment.page_used -= 1;
            }
            if segment.is_empty() {
                self.drop_segment(self.segment_group.leave(self, segment_ptr));
            }
        } else {
            if !page.is_in_delayed() {
                if page.set_in_delayed(true) {
                    // 当前线程修改竟态内存成功
                    self.delayed_week_up_pages.push(self, page_ptr);
                };
            }
        }
    }
}

impl AccessGuard<Segment> for Heap {
    fn before_access(&self, ptr: NonNull<Segment>) {
        debug_assert_eq!(ptr.as_ptr() as usize % SEGMENT_SIZE, 0)
    }
}

impl AccessGuard<Page> for Heap {
    fn before_access(&self, _ptr: NonNull<Page>) {
        // DO NOTHING
        // access 函数中调用了其他的 Guard，该做的检查其他的 Guard 会做
    }

    fn access<N>(&self, ptr: N) -> &mut Page
    where
        N: Into<NonNull<Page>>,
    {
        let ptr = ptr.into();
        let segment_ptr = Self::ptr_segment(ptr.cast().as_ptr())
            .expect("Never happened, the page ptr will not be null");
        let segment = self.access(segment_ptr);
        segment.access(ptr)
    }
}

#[repr(transparent)]
pub(crate) struct Heap(OsBox<HeapInner>);
impl Deref for Heap {
    type Target = HeapInner;

    fn deref(&self) -> &Self::Target {
        self.0.as_ref()
    }
}

impl DerefMut for Heap {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.0.as_mut()
    }
}

#[cfg(test)]
mod tests {
    use crate::consts::{BIN_BLOCK_SIZE, PTR_SIZE, SEGMENT_SIZE};
    use crate::heap::Heap;
    use crate::utils::access_guard::AccessGuard;
    use std::alloc::Layout;
    use std::iter;

    #[test]
    fn test_heap_basic() {
        const TEST_BIN: usize = 24;
        const TEST_ALLOC_BYTES: usize = BIN_BLOCK_SIZE[TEST_BIN] * PTR_SIZE;

        // test basic alloc/dealloc
        Heap::thread_local_with(|heap| {
            let layout = Layout::from_size_align(TEST_ALLOC_BYTES, TEST_ALLOC_BYTES).unwrap();
            assert!(heap.active_pages[TEST_BIN].first().is_none());
            let ptr = heap.alloc(layout);
            assert!(ptr.is_some());
            assert!(heap.active_pages[TEST_BIN].first().is_some());
            heap.dealloc(ptr.unwrap());
            assert!(heap.active_pages[TEST_BIN].first().is_none());
        });

        // 写满一个 page 测试 segment 状态

        // 大致的 segment 容量，没有计算各种 header 数据的占用
        let block_per_seg = (SEGMENT_SIZE + TEST_ALLOC_BYTES - 1) / TEST_ALLOC_BYTES;

        Heap::thread_local_with(|heap| {
            let mut ptr_vec = vec![];
            iter::repeat(()).take(block_per_seg).for_each(|_| {
                let layout = Layout::from_size_align(TEST_ALLOC_BYTES, TEST_ALLOC_BYTES).unwrap();
                let ptr = heap.alloc(layout);
                assert!(ptr.is_some());
                ptr_vec.push(ptr.unwrap());
            });
            assert!(heap.active_pages[TEST_BIN].first().is_some());
            assert!(heap.sleep_pages.first().is_some());
            let sleep_page = heap.access(heap.sleep_pages.first().unwrap());
            assert!(sleep_page.is_in_sleep());
            // segment_group 中应该由俩 segment
            assert_eq!(heap.segment_group.len(), 2);

            // 释放一半的 block
            let ptr_vec2 = ptr_vec.split_off(block_per_seg / 2);
            ptr_vec2.into_iter().for_each(|ptr| {
                heap.dealloc(ptr);
            });
            assert_eq!(heap.segment_group.len(), 1);
            assert!(heap.active_pages[TEST_BIN].first().is_some());
            assert!(heap.sleep_pages.first().is_some());

            // 释放全部的 block
            ptr_vec.into_iter().for_each(|ptr| {
                heap.dealloc(ptr);
            });
            assert_eq!(heap.segment_group.len(), 0);
            assert!(heap.active_pages[TEST_BIN].first().is_none());
            assert!(heap.sleep_pages.first().is_none());
        });
    }
}
